package com.zengde.ai_sandbox.service.impl;

import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChunk;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import io.reactivex.Flowable;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Service
@RequiredArgsConstructor
public class MyDbModelService {

    private final ArkService arkService;

    @Value("${my.doubao.think.model.name}")
    private String thinkModelName;

    @Value("${my.doubao.normal.model.name}")
    private String normalModelName;

    // 在服务类中添加
//    private final RateLimiter rateLimiter = RateLimiter.create(10); // 每秒10个请求


    // 存储多轮对话历史: sessionId -> 消息列表
    private final Map<String, List<ChatMessage>> conversationHistory = new ConcurrentHashMap<>();

    /**
     * 执行豆包异步流式调用
     * @param sessionId 会话ID
     * @param userInput 用户输入
     * @return Flux流式响应
     */
    public Flux<String> streamDoubaoResponse(String sessionId, String userInput) {
        // 获取或初始化对话历史
        List<ChatMessage> messages = conversationHistory.computeIfAbsent(sessionId,
                k -> new ArrayList<>(Collections.singletonList(
                        ChatMessage.builder()
                                .role(ChatMessageRole.SYSTEM)
                                .content("你是豆包，是由字节跳动开发的 AI 人工智能助手")
                                .build()
                ))
        );

        // 添加用户新消息
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(userInput)
                .build());

        // 构建请求
        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(thinkModelName) // 使用思考模型
                .messages(messages)
                .build();

        // 执行异步流式调用
        Flowable<ChatCompletionChunk> flowable = arkService.streamChatCompletion(request);

        // 转换为Spring Flux
        return Flux.create(fluxSink -> {
            StringBuilder fullResponse = new StringBuilder();

            flowable.subscribe(
                    response -> {
                        if (!response.getChoices().isEmpty()) {
                            String content = String.valueOf(response.getChoices().get(0).getMessage().getContent());
                            if (content != null) {

                                // 1. 确保内容为UTF-8字符串
                                String cleanContent = new String(
                                        content.getBytes(StandardCharsets.ISO_8859_1),
                                        StandardCharsets.UTF_8
                                );

                                // 发送数据块
                                fluxSink.next(cleanContent);
                                // 累积完整响应
                                fullResponse.append(cleanContent);
                            }
                        }
                    },
                    error -> {
                        log.error("豆包调用失败", error);
                        fluxSink.error(error);
                    },
                    () -> {
                        // 流结束时添加助手响应到历史
                        messages.add(ChatMessage.builder()
                                .role(ChatMessageRole.ASSISTANT)
                                .content(fullResponse.toString())
                                .build());

                        log.info("会话 {} 完成，历史记录: {} 条", sessionId, messages.size());
                        fluxSink.complete();
                    }
            );
        });
    }

    // 在MyDbModelService中添加
    public void resetSession(String sessionId) {
        conversationHistory.remove(sessionId);
        log.info("会话 {} 已重置", sessionId);
    }
}